/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: hccl实现基类头文件
 * Author: hanwei
 * Create: 2020-01-05
 */

#ifndef HCCL_IMPL_RPC_H
#define HCCL_IMPL_RPC_H

#include "hccl_impl_base.h"
#include "profiling_manager.h"

namespace hccl {

class HcclImplRpc : public HcclImplBase {
public:
    explicit HcclImplRpc();
    ~HcclImplRpc() override;

    HcclResult RegRemoteLookUpExecutor(s32 tag, bool isWorker, bool isKernelLaunch, HcclOpPtr opPtr = nullptr);

    HcclResult RegRemoteUpdateExecutor(s32 tag, bool isWorker, bool isKernelLaunch);

    HcclResult CollRecvLookupRequest(HcclOpPtr opPtr);

    HcclResult CollIsendLookupResponse(HcclOpPtr opPtr);

    HcclResult CollRecvUpdateRequest(HcclOpPtr opPtr);

    HcclResult CollRemoteLookUp(HcclOpPtr opPtr);
    HcclResult CollRemoteLookUpPaired(HcclOpPtr opPtr);
    HcclResult CollRemoteUpdate(HcclOpPtr opPtr);
    HcclResult CollRemoteUpdatePaired(HcclOpPtr opPtr);

    HcclResult WaitLookUpData(HcclOpPtr opPtr);

    HcclResult GetLookupRequest(HcclOpPtr opPtr);

    HcclResult IsetLookupResponse(HcclOpPtr opPtr);

    HcclResult RemoteLookUp(HcclOpPtr opPtr);

    HcclResult IsendUpdateResponse(HcclOpPtr opPtr);

    HcclResult CollRemoteUpdateKernelLaunch(HdcsEmbeddingServiceParam &para, void *stream);

    HcclResult RegTransportsKernelLaunch(s32 tag, const std::map<u32, u32> &psRankIdMap);

    HcclResult AbortSelf(s32 tag);
    HcclResult RegTransportLinks(s32 linkNum, void *transportPara);
    bool WorkerGetBoardVersion();
    bool PsGetBoardVersion();
    HcclResult WorkerGetBoardVersion(bool &isBoardVersion);

    HcclResult WaitSome(HcclOpPtr opPtr);
    HcclResult BuildHeterogeneousTransport(u32 commId, u32 peerRank, s32 tag, TransportHandle &transportHandle);
    HcclResult Isend(void* buffer, u64 count, HcclDataType dataType, u32 peerRank, s32 tag,
        HcclRequest &requestHandle, HcclUserRequire &userRequire);
    HcclResult Imrecv(void* buffer, u64 count, HcclDataType dataType, HcclMessage msgHandle,
        HcclRequest &requestHandle);
    HcclResult Improbe(u32 peerRank, s32 tag, s32 &flag, HcclMessage &msgHandle, HcclStatus &status);
    HcclResult HcclTest(HcclRequest requestHandle, s32 &flag, HcclStatus &compState);
    HcclResult NotifyAndResetKernelLaunch(const HdcsEmbeddingServiceParam &para, void *stream);
    HcclResult ValueReduceSumAicoreKernelLaunch(const HdcsEmbeddingServiceParam &para, void *stream);
    HcclResult ValueClearAicoreKernelLaunch(const HdcsEmbeddingServiceParam &para, void *stream);
    HcclResult UpdateTaskSendByRdma(const HdcsEmbeddingServiceParam &para, void *stream);

private:
    HcclResult SetupEsUpdatePara(EsUpdatePara &esUpdatePara, const HdcsEmbeddingServiceParam &para);
};
}
#endif